{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "20549885",
   "metadata": {},
   "source": [
    "# 使用 Milvus 和 DeepSeek 构建 RAG\n",
    "\n",
    "DeepSeek 帮助开发者使用高性能语言模型构建和扩展 AI 应用。它提供高效的推理、灵活的 API 以及先进的专家混合 (MoE) 架构，用于强大的推理和检索任务。\n",
    "\n",
    "在本教程中，我们将展示如何使用 Milvus 和 DeepSeek 构建一个检索增强生成 (RAG) 管道。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7394c701",
   "metadata": {},
   "source": [
    "## 准备工作"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89103a1e",
   "metadata": {},
   "source": [
    "### 依赖与环境"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9c18d7b4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: openai in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (1.82.0)\n",
      "Requirement already satisfied: requests in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (2.32.3)\n",
      "Requirement already satisfied: tqdm in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (4.67.1)\n",
      "Requirement already satisfied: torch in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (2.7.0)\n",
      "Requirement already satisfied: pymilvus[model] in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (2.5.10)\n",
      "Requirement already satisfied: setuptools>69 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (78.1.1)\n",
      "Requirement already satisfied: grpcio<=1.67.1,>=1.49.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (1.67.1)\n",
      "Requirement already satisfied: protobuf>=3.20.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (6.31.0)\n",
      "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (1.1.0)\n",
      "Requirement already satisfied: ujson>=2.0.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (5.10.0)\n",
      "Requirement already satisfied: pandas>=1.2.4 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (2.2.3)\n",
      "Requirement already satisfied: milvus-lite>=2.4.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (2.4.12)\n",
      "Requirement already satisfied: pymilvus.model>=0.3.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus[model]) (0.3.2)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (4.9.0)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (1.9.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (0.28.1)\n",
      "Requirement already satisfied: jiter<1,>=0.4.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (0.10.0)\n",
      "Requirement already satisfied: pydantic<3,>=1.9.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (2.11.4)\n",
      "Requirement already satisfied: sniffio in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (1.3.1)\n",
      "Requirement already satisfied: typing-extensions<5,>=4.11 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from openai) (4.13.2)\n",
      "Requirement already satisfied: idna>=2.8 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from anyio<5,>=3.5.0->openai) (3.10)\n",
      "Requirement already satisfied: certifi in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from httpx<1,>=0.23.0->openai) (2025.4.26)\n",
      "Requirement already satisfied: httpcore==1.* in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from httpx<1,>=0.23.0->openai) (1.0.9)\n",
      "Requirement already satisfied: h11>=0.16 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai) (0.16.0)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pydantic<3,>=1.9.0->openai) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.33.2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pydantic<3,>=1.9.0->openai) (2.33.2)\n",
      "Requirement already satisfied: typing-inspection>=0.4.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pydantic<3,>=1.9.0->openai) (0.4.0)\n",
      "Requirement already satisfied: charset_normalizer<4,>=2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from requests) (3.4.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from requests) (2.4.0)\n",
      "Requirement already satisfied: filelock in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (3.18.0)\n",
      "Requirement already satisfied: sympy>=1.13.3 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (1.14.0)\n",
      "Requirement already satisfied: networkx in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (3.1.6)\n",
      "Requirement already satisfied: fsspec in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (2025.5.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.6.77)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.6.77)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.6.80)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.5.1.17 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (9.5.1.17)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.6.4.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (11.3.0.4)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (10.3.7.77)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (11.7.1.2)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.5.4.2)\n",
      "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (0.6.3)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.26.2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (2.26.2)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.6.77)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (12.6.85)\n",
      "Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (1.11.1.6)\n",
      "Requirement already satisfied: triton==3.3.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from torch) (3.3.0)\n",
      "Requirement already satisfied: numpy>=1.26.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pandas>=1.2.4->pymilvus[model]) (2.2.6)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pandas>=1.2.4->pymilvus[model]) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pandas>=1.2.4->pymilvus[model]) (2025.2)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pandas>=1.2.4->pymilvus[model]) (2025.2)\n",
      "Requirement already satisfied: transformers>=4.36.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus.model>=0.3.0->pymilvus[model]) (4.52.2)\n",
      "Requirement already satisfied: onnxruntime in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus.model>=0.3.0->pymilvus[model]) (1.22.0)\n",
      "Requirement already satisfied: scipy>=1.10.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from pymilvus.model>=0.3.0->pymilvus[model]) (1.15.3)\n",
      "Requirement already satisfied: six>=1.5 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from python-dateutil>=2.8.2->pandas>=1.2.4->pymilvus[model]) (1.17.0)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from sympy>=1.13.3->torch) (1.3.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.30.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from transformers>=4.36.0->pymilvus.model>=0.3.0->pymilvus[model]) (0.31.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from transformers>=4.36.0->pymilvus.model>=0.3.0->pymilvus[model]) (25.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from transformers>=4.36.0->pymilvus.model>=0.3.0->pymilvus[model]) (6.0.2)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from transformers>=4.36.0->pymilvus.model>=0.3.0->pymilvus[model]) (2024.11.6)\n",
      "Requirement already satisfied: tokenizers<0.22,>=0.21 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from transformers>=4.36.0->pymilvus.model>=0.3.0->pymilvus[model]) (0.21.1)\n",
      "Requirement already satisfied: safetensors>=0.4.3 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from transformers>=4.36.0->pymilvus.model>=0.3.0->pymilvus[model]) (0.5.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from jinja2->torch) (3.0.2)\n",
      "Requirement already satisfied: coloredlogs in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from onnxruntime->pymilvus.model>=0.3.0->pymilvus[model]) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from onnxruntime->pymilvus.model>=0.3.0->pymilvus[model]) (25.2.10)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in /home/django/miniconda3/envs/deepseek/lib/python3.13/site-packages (from coloredlogs->onnxruntime->pymilvus.model>=0.3.0->pymilvus[model]) (10.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install \"pymilvus[model]==2.5.10\" openai==1.82.0 requests==2.32.3 tqdm==4.67.1 torch==2.7.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3c0999-d670-41a9-afbd-d8a020fe1631",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "375ad823",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# 从环境变量获取 DeepSeek API Key\n",
    "api_key = os.getenv(\"DEEPSEEK_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "21a19f48-d548-4b0c-a7c8-4d1e3f972a08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hugging face 镜像, 如果pymilvus还是连接超时，执行pip install ipywidgets --upgrade，升级 ipywidgets\n",
    "os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'\n",
    "os.environ['HUGGINGFACE_CO_URL_HOME'] = 'https://hf-mirror.com'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db44bb26",
   "metadata": {},
   "source": [
    "### 准备数据"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25f92a95",
   "metadata": {},
   "source": [
    "我们使用 Milvus 文档 2.4.x 中的 FAQ 页面作为我们 RAG 中的私有知识库，这是一个简单 RAG 管道的良好数据源。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1a8b9e2",
   "metadata": {},
   "source": [
    "下载 zip 文件并将文档解压到 `milvus_docs` 文件夹。\n",
    "\n",
    "**建议在命令行执行下面命令**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a81fa031",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!wget https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip\n",
    "#!unzip -q milvus_docs_2.4.x_en.zip -d milvus_docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1198466",
   "metadata": {},
   "source": [
    "我们从 `milvus_docs/en/faq` 文件夹加载所有 markdown 文件。对于每个文档，我们简单地使用 \"# \" 来分割文件中的内容，这样可以大致分离出 markdown 文件中每个主要部分的内容。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "c9035a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "\n",
    "text_lines = []\n",
    "\n",
    "for file_path in glob(\"milvus_docs/en/faq/*.md\", recursive=True):\n",
    "    with open(file_path, \"r\") as file:\n",
    "        file_text = file.read()\n",
    "\n",
    "    text_lines += file_text.split(\"# \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "01b73e74-ee7d-4daf-b7db-1c7a10bfc0bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "72"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(text_lines)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cc2a0b8",
   "metadata": {},
   "source": [
    "### 准备 LLM 和 Embedding 模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19eaff7a",
   "metadata": {},
   "source": [
    "DeepSeek 支持 OpenAI 风格的 API，您可以使用相同的 API 进行微小调整来调用 LLM。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b994eb47",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "deepseek_client = OpenAI(\n",
    "    api_key=api_key,\n",
    "    base_url=\"https://api.deepseek.com/v1\",  # DeepSeek API 的基地址\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cc5a5e2",
   "metadata": {},
   "source": [
    "定义一个 embedding 模型，使用 `milvus_model` 来生成文本嵌入。我们以 `DefaultEmbeddingFunction` 模型为例，这是一个预训练的轻量级嵌入模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4d7a0391-73b6-4dfd-8000-bce44747ebb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from pymilvus import (\n",
    "#     connections,      # 连接服务器\n",
    "#     Collection,       # 操作集合\n",
    "#     FieldSchema,      # 定义字段\n",
    "#     CollectionSchema, # 定义集合结构\n",
    "#     DataType          # 数据类型\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "3a94242a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import model as milvus_model\n",
    "\n",
    "embedding_model = milvus_model.DefaultEmbeddingFunction()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93fb1696",
   "metadata": {},
   "source": [
    "生成一个测试嵌入并打印其维度和前几个元素。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "88a27567",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "768\n",
      "[-0.04836059  0.07163021 -0.01130063 -0.03789341 -0.03320651 -0.01318453\n",
      " -0.03041721 -0.02269495 -0.02317858 -0.00426026]\n"
     ]
    }
   ],
   "source": [
    "test_embedding = embedding_model.encode_queries([\"This is a test\"])[0]\n",
    "embedding_dim = len(test_embedding)\n",
    "print(embedding_dim)\n",
    "print(test_embedding[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "f7683f3a-d9e4-4c8e-9a66-c341911bef6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.02752976  0.0608853   0.00388525 -0.00215193 -0.02774976 -0.0118618\n",
      " -0.04020916 -0.06023417 -0.03813156  0.0100272 ]\n"
     ]
    }
   ],
   "source": [
    "test_embedding_0 = embedding_model.encode_queries([\"That is a test\"])[0]\n",
    "print(test_embedding_0[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a778887",
   "metadata": {},
   "source": [
    "## 将数据加载到 Milvus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02b23a24",
   "metadata": {},
   "source": [
    "### 创建 Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6976a001-996f-44f7-82d0-e92f2af3479f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Milvus Client\n"
     ]
    }
   ],
   "source": [
    "# from pymilvus import MilvusClient\n",
    "# print(MilvusClient.__doc__)  # 应输出类文档说明"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "95e84b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(uri=\"./milvus_demo.db\")\n",
    "\n",
    "collection_name = \"my_rag_collection\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68648561",
   "metadata": {},
   "source": [
    "关于 `MilvusClient` 的参数：\n",
    "\n",
    "*   将 `uri` 设置为本地文件，例如 `./milvus.db`，是最方便的方法，因为它会自动利用 Milvus Lite 将所有数据存储在此文件中。\n",
    "*   如果您有大规模数据，可以在 Docker 或 Kubernetes 上设置性能更高的 Milvus 服务器。在此设置中，请使用服务器 URI，例如 `http://localhost:19530`，作为您的 `uri`。\n",
    "*   如果您想使用 Zilliz Cloud（Milvus 的完全托管云服务），请调整 `uri` 和 `token`，它们对应 Zilliz Cloud 中的 Public Endpoint 和 Api key。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ce1bf3e",
   "metadata": {},
   "source": [
    "检查 collection 是否已存在，如果存在则删除它。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "aee85c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "if milvus_client.has_collection(collection_name):\n",
    "    milvus_client.drop_collection(collection_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73eb379f",
   "metadata": {},
   "source": [
    "创建一个具有指定参数的新 collection。\n",
    "\n",
    "如果我们不指定任何字段信息，Milvus 将自动创建一个默认的 `id` 字段作为主键，以及一个 `vector` 字段来存储向量数据。一个保留的 JSON 字段用于存储非 schema 定义的字段及其值。\n",
    "\n",
    "`metric_type` (距离度量类型):\n",
    "     作用：定义如何计算向量之间的相似程度。\n",
    "     例如：`IP` (内积) - 值越大通常越相似；`L2` (欧氏距离) - 值越小越相似；`COSINE` (余弦相似度) - 通常转换为距离，值越小越相似。\n",
    "     选择依据：根据你的嵌入模型的特性和期望的相似性定义来选择。\n",
    "\n",
    " `consistency_level` (一致性级别):\n",
    "     作用：定义数据写入后，读取操作能多快看到这些新数据。\n",
    "     例如：\n",
    "         `Strong` (强一致性): 总是读到最新数据，可能稍慢。\n",
    "         `Bounded` (有界过期): 可能读到几秒内旧数据，性能较好 (默认)。\n",
    "         `Session` (会话一致性): 自己写入的自己能立刻读到。\n",
    "         `Eventually` (最终一致性): 最终会读到新数据，但没时间保证，性能最好。\n",
    "     选择依据：在数据实时性要求和系统性能之间做权衡。\n",
    "\n",
    "简单来说：\n",
    " `metric_type`：怎么算相似。\n",
    " `consistency_level`：新数据多久能被读到。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "bd0b2df8",
   "metadata": {},
   "outputs": [],
   "source": [
    "milvus_client.create_collection(\n",
    "    collection_name=collection_name,\n",
    "    dimension=embedding_dim,\n",
    "    metric_type=\"IP\",  # 内积距离\n",
    "    consistency_level=\"Strong\",  # 支持的值为 (`\"Strong\"`, `\"Session\"`, `\"Bounded\"`, `\"Eventually\"`)。更多详情请参见 https://milvus.io/docs/consistency.md#Consistency-Level。\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c15bafb",
   "metadata": {},
   "source": [
    "### 插入数据"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "171d3b35",
   "metadata": {},
   "source": [
    "遍历文本行，创建嵌入，然后将数据插入 Milvus。\n",
    "\n",
    "这里有一个新字段 `text`，它是在 collection schema 中未定义的字段。它将自动添加到保留的 JSON 动态字段中，该字段在高级别上可以被视为普通字段。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "ad077094",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating embeddings: 100%|██████████████████████████████████████████████████████████| 72/72 [00:00<00:00, 652245.98it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'insert_count': 72, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], 'cost': 0}"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "data = []\n",
    "\n",
    "doc_embeddings = embedding_model.encode_documents(text_lines)\n",
    "\n",
    "for i, line in enumerate(tqdm(text_lines, desc=\"Creating embeddings\")):\n",
    "    data.append({\"id\": i, \"vector\": doc_embeddings[i], \"text\": line})\n",
    "\n",
    "milvus_client.insert(collection_name=collection_name, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd971f6b",
   "metadata": {},
   "source": [
    "## 构建 RAG"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "534dc076",
   "metadata": {},
   "source": [
    "### 检索查询数据"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a6fd7e7",
   "metadata": {},
   "source": [
    "我们指定一个关于 Milvus 的常见问题。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "6e2f5c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"What factors impact CPU usage?\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52401a38",
   "metadata": {},
   "source": [
    "在 collection 中搜索该问题，并检索语义上最匹配的前3个结果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "0dd4cbac",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_res = milvus_client.search(\n",
    "    collection_name=collection_name,\n",
    "    data=embedding_model.encode_queries(\n",
    "        [question]\n",
    "    ),  # 将问题转换为嵌入向量\n",
    "    limit=3,  # 返回前3个结果\n",
    "    search_params={\"metric_type\": \"IP\", \"params\": {}},  # 内积距离\n",
    "    output_fields=[\"text\"],  # 返回 text 字段\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffcce135",
   "metadata": {},
   "source": [
    "让我们看一下查询的搜索结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "6a7f6eb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "    [\n",
      "        \"What factors impact CPU usage?\\n\\nCPU usage increases when Milvus is building indexes or running queries. In general, index building is CPU intensive except when using Annoy, which runs on a single thread.\\n\\nWhen running queries, CPU usage is affected by `nq` and `nprobe`. When `nq` and `nprobe` are small, concurrency is low and CPU usage stays low.\\n\\n###\",\n",
      "        0.7029333710670471\n",
      "    ],\n",
      "    [\n",
      "        \"How do I know if my CPU supports Milvus?\\n\\n{{fragments/cpu_support.md}}\\n\\n###\",\n",
      "        0.4407930076122284\n",
      "    ],\n",
      "    [\n",
      "        \"Does simultaneously inserting data and searching impact query performance?\\n\\nInsert operations are not CPU intensive. However, because new segments may not have reached the threshold for index building, Milvus resorts to brute-force search\\u2014significantly impacting query performance.\\n\\nThe `rootcoord.minSegmentSizeToEnableIndex` parameter determines the index-building threshold for a segment, and is set to 1024 rows by default. See [System Configuration](system_configuration.md) for more information.\\n\\n###\",\n",
      "        0.3326629400253296\n",
      "    ]\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "retrieved_lines_with_distances = [\n",
    "    (res[\"entity\"][\"text\"], res[\"distance\"]) for res in search_res[0]\n",
    "]\n",
    "print(json.dumps(retrieved_lines_with_distances, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd4c186",
   "metadata": {},
   "source": [
    "### 使用 LLM 获取 RAG 响应"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cd1ae3a",
   "metadata": {},
   "source": [
    "将检索到的文档转换为字符串格式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "0676448f",
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"\\n\".join(\n",
    "    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "107df42a-b3f7-48a8-b66b-fc82fe3ec174",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'What factors impact CPU usage?\\n\\nCPU usage increases when Milvus is building indexes or running queries. In general, index building is CPU intensive except when using Annoy, which runs on a single thread.\\n\\nWhen running queries, CPU usage is affected by `nq` and `nprobe`. When `nq` and `nprobe` are small, concurrency is low and CPU usage stays low.\\n\\n###\\nHow do I know if my CPU supports Milvus?\\n\\n{{fragments/cpu_support.md}}\\n\\n###\\nDoes simultaneously inserting data and searching impact query performance?\\n\\nInsert operations are not CPU intensive. However, because new segments may not have reached the threshold for index building, Milvus resorts to brute-force search—significantly impacting query performance.\\n\\nThe `rootcoord.minSegmentSizeToEnableIndex` parameter determines the index-building threshold for a segment, and is set to 1024 rows by default. See [System Configuration](system_configuration.md) for more information.\\n\\n###'"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "19998758-7f98-4cb8-8789-625fcfaad00e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'What factors impact CPU usage?'"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ad25756",
   "metadata": {},
   "source": [
    "为语言模型定义系统和用户提示。此提示是使用从 Milvus 检索到的文档组装而成的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "b655f6f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = \"\"\"\n",
    "Human: 你是一个 AI 助手。你能够从提供的上下文段落片段中找到问题的答案。\n",
    "\"\"\"\n",
    "USER_PROMPT = f\"\"\"\n",
    "请使用以下用 <context> 标签括起来的信息片段来回答用 <question> 标签括起来的问题。最后追加原始回答的中文翻译，并用 <translated>和</translated> 标签标注。\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "<translated>\n",
    "</translated>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "97089c31-f85c-47a9-8498-78520513bc67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n请使用以下用 <context> 标签括起来的信息片段来回答用 <question> 标签括起来的问题。最后追加原始回答的中文翻译，并用 <translated>和</translated> 标签标注。\\n<context>\\nWhat factors impact CPU usage?\\n\\nCPU usage increases when Milvus is building indexes or running queries. In general, index building is CPU intensive except when using Annoy, which runs on a single thread.\\n\\nWhen running queries, CPU usage is affected by `nq` and `nprobe`. When `nq` and `nprobe` are small, concurrency is low and CPU usage stays low.\\n\\n###\\nHow do I know if my CPU supports Milvus?\\n\\n{{fragments/cpu_support.md}}\\n\\n###\\nDoes simultaneously inserting data and searching impact query performance?\\n\\nInsert operations are not CPU intensive. However, because new segments may not have reached the threshold for index building, Milvus resorts to brute-force search—significantly impacting query performance.\\n\\nThe `rootcoord.minSegmentSizeToEnableIndex` parameter determines the index-building threshold for a segment, and is set to 1024 rows by default. See [System Configuration](system_configuration.md) for more information.\\n\\n###\\n</context>\\n<question>\\nWhat factors impact CPU usage?\\n</question>\\n<translated>\\n</translated>\\n'"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "USER_PROMPT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "184b457f",
   "metadata": {},
   "source": [
    "使用 DeepSeek 提供的 `deepseek-chat` 模型根据提示生成响应。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "638a7561",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The factors that impact CPU usage are:\n",
      "1. **Index Building**: CPU usage increases when Milvus is building indexes (except when using Annoy, which runs on a single thread).\n",
      "2. **Query Execution**: CPU usage is affected by the parameters `nq` (number of queries) and `nprobe` (number of clusters to search). When these values are small, concurrency is low, and CPU usage remains low.\n",
      "\n",
      "<translated>\n",
      "影响CPU使用率的因素包括：\n",
      "1. **索引构建**：当Milvus构建索引时，CPU使用率会增加（使用Annoy索引时除外，因为Annoy在单线程上运行）。\n",
      "2. **查询执行**：CPU使用率受参数`nq`（查询数量）和`nprobe`（搜索的聚类数量）影响。当这些值较小时，并发性较低，CPU使用率也会保持在较低水平。\n",
      "</translated>\n"
     ]
    }
   ],
   "source": [
    "response = deepseek_client.chat.completions.create(\n",
    "    model=\"deepseek-chat\",\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": USER_PROMPT},\n",
    "    ],\n",
    ")\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a1a6cd5-f94b-4df5-b1be-c2827b643b25",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b6e893b8-154f-4a03-8b7e-e1fc3f9c3a83",
   "metadata": {},
   "source": [
    "## 读取民法典示例数据 mfd.md"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c660c70e-04f5-4289-88bd-c12868119f67",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_clauses(markdown_text: str) -> list[tuple[str, str]]:\n",
    "    \"\"\"\n",
    "    从Markdown文本提取法律条款\n",
    "    \n",
    "    返回: [(条款编号, 内容), ...]\n",
    "    \"\"\"\n",
    "    try:\n",
    "        pattern = re.compile(r'\\*\\*(.*?)\\*\\*\\s*(.*?)(?=\\n\\*\\*|\\Z)', re.DOTALL)\n",
    "        return [num + '：' + content.strip() for num, content in pattern.findall(markdown_text)]\n",
    "    except Exception as e:\n",
    "        print(f\"解析失败: {e}\")\n",
    "        return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "50ce3b55-7daa-4727-95bd-3b00f6aa2b03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import re\n",
    "\n",
    "law_lines = []\n",
    "\n",
    "for file_path in glob(\"mfd.md\", recursive=True):\n",
    "    with open(file_path, \"r\") as file:\n",
    "        file_text = file.read()\n",
    "    law_lines = extract_clauses(file_text)\n",
    "    \n",
    "    # for i, (num, content) in enumerate(contens, 1):\n",
    "        # print(f\"{i}. {num}: {content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "44f06e59-f2ac-4e6c-aeb3-9998fd1690d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "387"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(law_lines)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da19441a-9508-4409-8e11-9d19f1c03534",
   "metadata": {},
   "source": [
    "## 准备LLM，上面已经生成的DeepSeek的实例 deepseek_client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b86c8d-21ae-4d76-b465-08b01a8efd98",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2c9a1359-6f8e-458f-ba5a-70c202817e31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "768\n",
      "[-0.04836059  0.07163021 -0.01130063 -0.03789341 -0.03320651 -0.01318453\n",
      " -0.03041721 -0.02269495 -0.02317858 -0.00426026]\n"
     ]
    }
   ],
   "source": [
    "test_embedding = embedding_model.encode_queries([\"This is a test\"])[0]\n",
    "embedding_dim = len(test_embedding)\n",
    "print(embedding_dim)\n",
    "print(test_embedding[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e362d793-7a93-4a41-baf4-430074ef4d7e",
   "metadata": {},
   "source": [
    "## 将数据加载到 Milvus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4252d0c-e42e-416f-b726-90d630d8731e",
   "metadata": {},
   "source": [
    "### 创建 Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ec3a5b4d-3dfd-45a2-b3aa-7418c348c775",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda/envs/py311/lib/python3.11/site-packages/milvus_lite/__init__.py:15: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
      "  from pkg_resources import DistributionNotFound, get_distribution\n"
     ]
    }
   ],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(uri=\"./milvus_demo.db\")\n",
    "\n",
    "collection_name = \"my_rag_collection_mfd\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "fd5a89d1-3a6f-47df-bf70-c9bf92401863",
   "metadata": {},
   "outputs": [],
   "source": [
    "if milvus_client.has_collection(collection_name):\n",
    "    milvus_client.drop_collection(collection_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "670ab9e5-a794-4823-8d40-5314a7962562",
   "metadata": {},
   "outputs": [],
   "source": [
    "milvus_client.create_collection(\n",
    "    collection_name=collection_name,\n",
    "    dimension=embedding_dim,\n",
    "    metric_type=\"IP\",  # 内积距离\n",
    "    consistency_level=\"Strong\",  # 支持的值为 (`\"Strong\"`, `\"Session\"`, `\"Bounded\"`, `\"Eventually\"`)。更多详情请参见 https://milvus.io/docs/consistency.md#Consistency-Level。\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e7047f6-fcc6-4c29-af4b-674ff0b752b7",
   "metadata": {},
   "source": [
    "### 插入数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "470b4e08-3c55-4255-a2cb-fb2c844967f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating embeddings: 100%|███████████████████████████████████████████████████████| 387/387 [00:00<00:00, 1712231.70it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'insert_count': 387, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384, 385, 386], 'cost': 0}"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "data = []\n",
    "\n",
    "doc_embeddings = embedding_model.encode_documents(law_lines)\n",
    "\n",
    "for i, line in enumerate(tqdm(law_lines, desc=\"Creating embeddings\")):\n",
    "    data.append({\"id\": i, \"vector\": doc_embeddings[i], \"text\": line})\n",
    "\n",
    "milvus_client.insert(collection_name=collection_name, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc4e567-08a6-40ac-8992-00215765c521",
   "metadata": {},
   "source": [
    "## 构建RAG"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f4f5c0-b8c3-4ccb-b5b2-6856921bfab7",
   "metadata": {},
   "source": [
    "### 检索数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "275c83d0-5d90-4589-b805-0a9a3bc2fd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"符合哪些情形合同的权利义务终止?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "ae1e7977-b965-474a-97c1-50733148b804",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_res = milvus_client.search(\n",
    "    collection_name=collection_name,\n",
    "    data=embedding_model.encode_queries(\n",
    "        [question]\n",
    "    ),  # 将问题转换为嵌入向量\n",
    "    limit=3,  # 返回前3个结果\n",
    "    search_params={\"metric_type\": \"IP\", \"params\": {}},  # 内积距离\n",
    "    output_fields=[\"text\"],  # 返回 text 字段\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "5f6c2467-2f4a-4650-a600-ec7e230bebc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "    [\n",
      "        \"\\u7b2c\\u56db\\u767e\\u56db\\u5341\\u4e5d\\u6761\\uff1a\\u53ef\\u4ee5\\u51fa\\u8d28\\u7684\\u6743\\u5229\\u5305\\u62ec\\uff1a\\n\\uff08\\u4e00\\uff09\\u6c47\\u7968\\u3001\\u672c\\u7968\\u3001\\u652f\\u7968\\uff1b\\n\\uff08\\u4e8c\\uff09\\u503a\\u5238\\u3001\\u5b58\\u6b3e\\u5355\\uff1b\\n\\uff08\\u4e09\\uff09\\u4ed3\\u5355\\u3001\\u63d0\\u5355\\uff1b\\n\\uff08\\u56db\\uff09\\u53ef\\u4ee5\\u8f6c\\u8ba9\\u7684\\u57fa\\u91d1\\u4efd\\u989d\\u3001\\u80a1\\u6743\\uff1b\\n\\uff08\\u4e94\\uff09\\u53ef\\u4ee5\\u8f6c\\u8ba9\\u7684\\u6ce8\\u518c\\u5546\\u6807\\u4e13\\u7528\\u6743\\u3001\\u4e13\\u5229\\u6743\\u3001\\u8457\\u4f5c\\u6743\\u7b49\\u77e5\\u8bc6\\u4ea7\\u6743\\u4e2d\\u7684\\u8d22\\u4ea7\\u6743\\uff1b\\n\\uff08\\u516d\\uff09\\u5e94\\u6536\\u8d26\\u6b3e\\uff1b\\n\\uff08\\u4e03\\uff09\\u6cd5\\u5f8b\\u3001\\u884c\\u653f\\u6cd5\\u89c4\\u89c4\\u5b9a\\u53ef\\u4ee5\\u51fa\\u8d28\\u7684\\u5176\\u4ed6\\u8d22\\u4ea7\\u6743\\u5229\\u3002\",\n",
      "        0.6444354057312012\n",
      "    ],\n",
      "    [\n",
      "        \"\\u7b2c\\u4e09\\u767e\\u4e00\\u5341\\u56db\\u6761\\uff1a\\u6309\\u4efd\\u5171\\u6709\\u4eba\\u5bf9\\u5171\\u6709\\u7684\\u4e0d\\u52a8\\u4ea7\\u6216\\u8005\\u52a8\\u4ea7\\u6309\\u7167\\u5176\\u4efd\\u989d\\u4eab\\u6709\\u6240\\u6709\\u6743\\u3002\",\n",
      "        0.6402398347854614\n",
      "    ],\n",
      "    [\n",
      "        \"\\u7b2c\\u4e8c\\u767e\\u4e03\\u5341\\u4e03\\u6761\\uff1a\\u4e1a\\u4e3b\\u4e0d\\u5f97\\u4ee5\\u653e\\u5f03\\u6743\\u5229\\u4e3a\\u7531\\u4e0d\\u5c65\\u884c\\u4e49\\u52a1\\u3002\",\n",
      "        0.6402398347854614\n",
      "    ]\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "retrieved_lawLines_with_distances = [\n",
    "    (res[\"entity\"][\"text\"], res[\"distance\"]) for res in search_res[0]\n",
    "]\n",
    "print(json.dumps(retrieved_lawLines_with_distances, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d937f61-7f62-469c-8789-0f3bf36a0f6f",
   "metadata": {},
   "source": [
    "### LLM 获取响应数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "b02e28ad-d305-44db-9dbe-398f16ddc3fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('第四百四十九条：可以出质的权利包括：\\n（一）汇票、本票、支票；\\n（二）债券、存款单；\\n（三）仓单、提单；\\n（四）可以转让的基金份额、股权；\\n（五）可以转让的注册商标专用权、专利权、著作权等知识产权中的财产权；\\n（六）应收账款；\\n（七）法律、行政法规规定可以出质的其他财产权利。',\n",
       " 0.6444354057312012)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retrieved_lawLines_with_distances[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "8500b38c-0c2f-45cd-a126-7abae26e1723",
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"\\n\".join(\n",
    "    [line_with_distance[0] for line_with_distance in retrieved_lawLines_with_distances]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "e5390070-2aff-48a3-ac57-4f49c833dd5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'第四百四十九条：可以出质的权利包括：\\n（一）汇票、本票、支票；\\n（二）债券、存款单；\\n（三）仓单、提单；\\n（四）可以转让的基金份额、股权；\\n（五）可以转让的注册商标专用权、专利权、著作权等知识产权中的财产权；\\n（六）应收账款；\\n（七）法律、行政法规规定可以出质的其他财产权利。\\n第三百一十四条：按份共有人对共有的不动产或者动产按照其份额享有所有权。\\n第二百七十七条：业主不得以放弃权利为由不履行义务。'"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "6c890e76-bcbe-445b-955f-9dd2734970fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = \"\"\"\n",
    "Human: 你是一个 AI 助手。你能够从提供的上下文段落片段中找到问题的答案。\n",
    "\"\"\"\n",
    "USER_PROMPT = f\"\"\"\n",
    "请使用以下用 <context> 标签括起来的信息片段来回答用 <question> 标签括起来的问题。\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a6ef349-f0c2-4829-887b-0ca46515f9c1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "0a021811-a4bb-4f3c-9e02-502e10fc5e31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "根据提供的上下文信息，该段落并未包含关于\"合同权利义务终止情形\"的相关法律规定。上下文主要涉及以下内容：\n",
      "1. 可以出质的权利类型（第四百四十九条）\n",
      "2. 按份共有人的所有权（第三百一十四条）\n",
      "3. 业主的义务（第二百七十七条）\n",
      "\n",
      "要回答合同权利义务终止的情形，需要参考《民法典》合同编中关于合同权利义务终止的具体条款（如第五百五十七条等），但这些内容并未包含在当前提供的上下文段落中。建议提供更相关的法律条文上下文以便准确回答。\n"
     ]
    }
   ],
   "source": [
    "response = deepseek_client.chat.completions.create(\n",
    "    model=\"deepseek-chat\",\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": USER_PROMPT},\n",
    "    ],\n",
    ")\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c777e093-7755-42b2-889e-219a7b3caf0e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
